#include <string_view>

#include "pqc.h"

int version, k, eta1, eta2, du, dv;
using namespace std::literals;

void generate_inequations(bool filtered, long long int sd, int numIneq_each)
{
    ////Alice generate and fix public and private keys
    // int sd = 10*time(NULL);
    //  long int sd = 16832015360;
    // generator.seed(0); //这里先固定随机数种子为0，方便复现实验
    // sd = 16832015360;
    generator.seed(sd);

    string fileprefix;

    // 文件路径
    if (filtered)
    {
        fileprefix = "bs_code/inequalities/filtered/cutfilter_Kyber" + to_string(version) + "Simplified_seed" + to_string(sd);
    }
    else
    {
        fileprefix = "bs_code/inequalities/unfiltered/unfilter_Kyber" + to_string(version) + "Simplified_seed" + to_string(sd);
    }

    // 获取最后一个斜杠的位置
    size_t last_slash_pos = fileprefix.find_last_of('/');

    // 如果存在斜杠，则提取斜杠之前的内容
    string folder = fileprefix.substr(0, last_slash_pos); 

    // 检查路径是否存在，如果不存在则创建它
    if (!fs::exists(folder))
    {
        cout << "Directory doesn't exist. Creating directory: " << folder << endl;
        if (!fs::create_directories(folder))
        {
            cerr << "Failed to create directory: " << folder << endl;
        }
    }
    else
    {
        cout << "Directory already exists: " << folder << endl;
    }

    ofstream myout_pkA(fileprefix + "_pkA.txt"); // file for public key A
    ofstream myout_pkt(fileprefix + "_pkt.txt"); // file for public key t
    ofstream myout_Coeff(fileprefix + ".txt");   // file for coefficient matrix
    ofstream myout_e_s(fileprefix + "_es.txt");  // file for solution
    ofstream myout_lvu(fileprefix + "_lvu.txt"); // file for lower bound, true value, upper bound

    // 检查文件是否成功打开
    if (!myout_pkA.is_open() || !myout_pkt.is_open() || !myout_Coeff.is_open() || !myout_e_s.is_open() || !myout_lvu.is_open())
    {
        cerr << "Failed to open files!" << endl;
        return;
    }

    // generate A from R_q
    vector<vector<vector<int>>> A(k, vector<vector<int>>(k, vector<int>(n)));
    vector<vector<int>> lwe_A(n * k, vector<int>(n * k));
    vector<vector<int>> temp(n, vector<int>(n));
    for (int i = 0; i < k; i++)
    {
        for (int j = 0; j < k; j++)
        {
            A[i][j] = sample_poly_U();
            temp = coeff2matrix(A[i][j]);
            for (int ii = 0; ii < n; ii++)
            {
                for (int jj = 0; jj < n; jj++)
                {
                    lwe_A[i * n + ii][j * n + jj] = temp[ii][jj];
                }
            }
        }
    }
    // store lwe A
    for (int i = 0; i < k * n; i++)
    {
        for (int j = 0; j < k * n; j++)
        {
            myout_pkA << lwe_A[i][j] << " ";
        }
        myout_pkA << endl;
    }

    // sample  s from B_eta1
    vector<vector<int>> s(k, vector<int>(n));
    for (int i = 0; i < k; i++)
    {
        s[i] = sample_poly_B(eta1);
    }
    // sample e from B_eta1;
    vector<vector<int>> e(k, vector<int>(n));
    vector<int> sum_e(n);
    for (int i = 0; i < k; i++)
    {
        e[i] = sample_poly_B(eta1);
    }

    // store (e[0]+..+e[k-1])
    for (int j = 0; j < n; j++)
    {
        sum_e[j] = 0;
        for (int i = 0; i < k; i++)
        {
            sum_e[j] += e[i][j];
        }
        myout_e_s << sum_e[j] << endl;
    }

    // store s
    for (int i = 0; i < k; i++)
    {
        for (int j = 0; j < n; j++)
        {
            myout_e_s << s[i][j] << endl;
        }
    }

    // calculate t =  As + e and store lwe_t
    vector<vector<int>> t(k, vector<int>(n));
    for (int i = 0; i < k; i++)
    {
        t[i] = dotproduct(A[i], s) + e[i];
    }
    for (int i = 0; i < k; i++)
    {
        for (int j = 0; j < n; j++)
        {
            myout_pkt << t[i][j] << endl;
        }
    }

    int numIneq = 0;
    int flag;
    int index = 0;     // v'的256个系数的索引
    int collected = 0; // 已经收集的方程个数

    vector<vector<int>> addN(n, vector<int>(numIneq_each)); // numIneq_each作为参数传入，代表每个index下收集多少个不等式。

    int lowerbound, upperbound; // 不等式的下界和上界
    int delta_N;                // v'走N步得到的d与正确d的差
    int delta_Nadd1;            // v'走（N+1）步得到的d与正确d的差

    vector<int> m(n);
    vector<int> m_prime(n);
    vector<int> a(n);
    vector<int> diff_v;                            // v'' - v 和主流论文中一致
    vector<vector<int>> diff_u(k, vector<int>(n)); // u'' - u 和主流论文中一致
    vector<vector<int>> Coeff_eq;                  // 收集的不等式的系数矩阵
    int d;
    while ((index < n) & (collected < numIneq_each))
    {
        // Bob sample a fixed m and all random values
        // sample m from {0,1}
        vector<int> m(n);
        for (int i = 0; i < n; i++)
        {
            m[i] = distribution(generator); // 这里是对m的真实赋值
        }

        // sample e2 from B_eta2
        vector<int> e2(n);
        e2 = sample_poly_B(eta2);

        // sample r from B_eta1 and r[0] = r[1] = .. = r[k]
        vector<vector<int>> r(k, vector<int>(n));
        r[0] = sample_poly_B(eta1);
        for (int i = 1; i < k; i++)
        {
            r[i] = r[0];
        }

        // calculate v = t'r + e2 +Decompress_q(m, 1)
        vector<int> v(n);
        v = dotproduct(t, r) + e2 + Decompress_q(m, 1);

        // calculate v' = Compress_q(v)
        vector<int> v_prime(n);
        v_prime = Compress_q(v, dv);

        // V''
        vector<int> v_primeprime(n);
        v_primeprime = Decompress_q(v_prime, dv);

        // Δv
        diff_v = biased_mod(plain_minus(v_primeprime, v), q); // 和主流论文一致

        // ciphertexts filter
        if (filtered)
        {
            if (abs(e2[index] + diff_v[index]) > 10)
            { // 和原版代码筛出的密文不同
                continue;
            }
        }

        // sample e1 from B_eta2
        vector<vector<int>> e1(k, vector<int>(n));
        for (int i = 0; i < k; i++)
        {
            e1[i] = sample_poly_B(eta2);
        }

        // calculate u = A'r + e1. Note: Tansform of A
        vector<vector<int>> u(k, vector<int>(n));
        for (int i = 0; i < k; i++)
        {
            vector<vector<int>> A_current_col(k, vector<int>(n));
            for (int j = 0; j < k; j++)
            {
                A_current_col[j] = A[j][i];
            }
            u[i] = dotproduct(A_current_col, r) + e1[i];
        }

        // calculate u' = Compress_q(u)
        vector<vector<int>> u_prime(k, vector<int>(n));
        for (int i = 0; i < k; i++)
        {
            u_prime[i] = Compress_q(u[i], du);
        }

        // calculate u'' = Decompress_q(u') = Decompress_q(Compress_q(u))
        vector<vector<int>> u_primeprime(k, vector<int>(n));
        for (int i = 0; i < k; i++)
        {
            u_primeprime[i] = Decompress_q(u_prime[i], du);
        }

        // Δu
        for (int i = 0; i < k; i++)
        {
            diff_u[i] = biased_mod(plain_minus(u_primeprime[i], u[i]), q); // 和主流论文一致
        }

        // -(Δu+e1)
        vector<vector<int>> neg_e1_plus_diffu(k, vector<int>(n));
        for (int i = 0; i < k; i++)
        {
            neg_e1_plus_diffu[i] = plain_minus(-1 * e1[i], diff_u[i]); // 作为已知项出现在d中，取原始值
        }

        // 主流论文中的误差 d
        d = plain_add(plain_add(plain_dotproduct(r, e), plain_dotproduct(neg_e1_plus_diffu, s)), plain_add(e2, diff_v))[index];

        // 诚实条件下的解密错误
        if ((d >= q / 4.0) | (d <= -q / 4.0))
        {
            printf("d = %d \n", d);
            for (int i = 0; i < k; i++)
            {
                print_vector(e1[i]);
            }

            printf("s * diff_u = %d\n", diff_v[index] - plain_dotproduct(diff_u, s)[index]);
            return;
        }

        vector<int> coeff_plain(n + k * n);
        coeff_plain = coeffe_coeffs(r[0], neg_e1_plus_diffu, index);

        // Bob alter v' step by step
        vector<int> v_prime_error(n);
        v_prime_error = v_prime;
        int v_primeprime_error;

        // 逐步增加步长获取不等式
        v_prime_error[index]++;
        v_primeprime_error = Decompress_q(mod(v_prime_error[index], pow(2, dv)), dv);
        m_prime[index] = Compress_q(mod(v_primeprime_error - dotproduct(s, u_primeprime)[index], q), 1);
        while (m[index] == m_prime[index])
        {
            v_prime_error[index]++;
            v_primeprime_error = Decompress_q(mod(v_prime_error[index], pow(2, dv)), dv);
            // calculate m'
            m_prime[index] = Compress_q(mod(v_primeprime_error - dotproduct(s, u_primeprime)[index], q), 1);
        };

        addN[index][collected] = v_prime_error[index] - v_prime[index] - 1;

        delta_N = Decompress_q(mod(v_prime[index] + addN[index][collected], pow(2, dv)), dv) - Decompress_q(v_prime[index], dv);
        delta_Nadd1 = Decompress_q(mod(v_prime[index] + addN[index][collected] + 1, pow(2, dv)), dv) - Decompress_q(v_prime[index], dv);
        // Gurobi: We do not support strict less-than, strict greater-than, or not-equal comparators
        //  Specifically, you can constrain an expression to be less-than-or-equal, greater-than-or-equal, or equal another.
        if (m[index] == 0)
        { // 0 -> 1
            lowerbound = biased_mod(ceil(q / 4.0) - delta_Nadd1, q);
            upperbound = biased_mod(floor(q / 4.0) - delta_N, q);
        }
        else
        { // 1 -> 0
            lowerbound = biased_mod(ceil(3 * q / 4.0) - round(q / 2.0) - delta_Nadd1, q);
            upperbound = biased_mod(floor(3 * q / 4.0) - round(q / 2.0) - delta_N, q);
        }

        // 验证不等式
        int value_plain;
        value_plain = 0;
        for (int j = 0; j < n; j++)
        {
            value_plain += coeff_plain[j] * sum_e[j];
        }
        for (int i = 0; i < k; i++)
        {
            for (int j = 0; j < n; j++)
            {
                value_plain += coeff_plain[n + i * n + j] * s[i][j];
            }
        }
        if ((value_plain <= upperbound - e2[index] - diff_v[index]) & (value_plain >= lowerbound - e2[index] - diff_v[index]))
        {
            myout_lvu << lowerbound - e2[index] - diff_v[index] << " " << value_plain << " " << upperbound - e2[index] - diff_v[index] << endl; //

            for (int j = 0; j < n; j++)
            {
                myout_Coeff << coeff_plain[j] << " "; // 256个 (e[0]+e[1]+..+e[k]) 的分量的系数
            }

            for (int i = n; i < n * (1 + k); i++)
            {
                myout_Coeff << coeff_plain[i] << " "; //
            }
            myout_Coeff << endl;
        }
        else
        {
            cout << "error value_plain = " << value_plain << ", bound = [" << lowerbound - e2[index] - diff_v[index] << ", " << upperbound - e2[index] - diff_v[index] << "], " << endl;
            cout << "d = " << d << " [" << lowerbound << ", " << upperbound << "]" << endl;
            // cout <<"current ofs_prem = " << ofs_prem << endl;

            // 即使错了也先放进去，看看最后有多少错的。
            // myout_lvu << lowerbound - e2[index] - diff_v[index] <<" "<< value_plain <<" "<< upperbound - e2[index] - diff_v[index]<<endl;//

            // return;
        }
        index++;
        if (index == n)
        {
            collected++;
            cout << "Collected " << collected << " inequalities for all indexes" << endl;
            index = 0;
        }
    }

    myout_pkA.close();
    myout_pkt.close();
    myout_Coeff.close();
    myout_e_s.close();
    myout_lvu.close();
}

int main(int argc, char *argv[])
{
    bool filtered;
    int numIneq_each;

    std::string s(argv[1]);
    if (s == "filtered")
    {
        filtered = true;
        printf("Generating filterd inequalities...");
        numIneq_each = 50; // 筛选的产生50 * 256 = 12800个方程
    }
    else if (s == "unfiltered")
    {
        filtered = false;
        printf("Generating unfilterd inequalities...");
        numIneq_each = 80; // 未筛选的产生80 * 256 = 20480个方程
    }
    else
    {
        printf("No indication of filtering. Give 'filtered' or 'unfiltered' \n");
        return 0;
    }

    long long int seed;
    seed = atoll(argv[2]);
    printf("Using seed %lld... ", seed);

    version = atoi(argv[3]);
    if (version == 512)
    {
        printf("Version 512.\n");
        k = 2;
        eta1 = 3;
        eta2 = 2;
        du = 10;
        dv = 4;
    }
    else if (version == 768)
    {
        printf("Version 768.\n");
        k = 3;
        eta1 = 2;
        eta2 == 2;
        du = 10;
        dv = 4;
    }
    else if (version == 1024)
    {
        printf("Version 1024. \n");
        k = 4;
        eta1 = 2;
        eta2 = 2;
        du = 11;
        dv = 5;
    }
    else
    {
        printf("No indication of version. Give '512', '768' or '1024' \n");
        return 0;
    }

    // test_diffIneq_validity();

    generate_inequations(filtered, seed, numIneq_each);

    return 0;
}
